#define _CRT_SECURE_NO_WARNINGS 1
#include <winsock2.h>
#include <ws2tcpip.h>
#include <iostream>
#include <cstring>
#include <ctime>

using namespace std;

constexpr int NTP_PACKET_SIZE = 48;
constexpr uint8_t NTP_HEADER_MODE_CLIENT = 3;
constexpr uint8_t NTP_HEADER_VERSION = 4;
constexpr uint64_t NTP_EPOCH_OFFSET = 2208988800ULL;

struct NTPPacket {
    uint8_t liVnMode{ 0x24 };  // LI=0, VN=4, Mode=4
    uint8_t stratum{ 2 };
    uint8_t poll{ 6 };
    uint8_t precision{ -6 };
    uint32_t rootDelay{ 0 };
    uint32_t rootDispersion{ 0 };
    uint32_t refId{ 'OOPC' };
    uint64_t refTime{ 0 };
    uint64_t origTime{ 0 };
    uint64_t recvTime{ 0 };
    uint64_t xmitTime{ 0 };
} __attribute__((packed));

int main(int argc, char* argv[]) 
{
    if (argc != 2) 
    {
        cerr << "Usage: " << argv[0] << " \"YYYY-MM-DD HH:MM:SS\"" << endl;
        return 1;
    }

    WSADATA wsa;
    if (WSAStartup(MAKEWORD(2, 2), &wsa) != 0) return 1;

    tm tmTime{};
    if (sscanf(argv[1], "%d-%d-%d %d:%d:%d",
        &tmTime.tm_year, &tmTime.tm_mon, &tmTime.tm_mday,
        &tmTime.tm_hour, &tmTime.tm_min, &tmTime.tm_sec) != 6) 
    {
        WSACleanup();
        return 1;
    }
    tmTime.tm_year -= 1900;
    tmTime.tm_mon--;
    tmTime.tm_isdst = -1;
    time_t unixTime = mktime(&tmTime);
    if (unixTime == -1) {
        WSACleanup();
        return 1;
    }
    uint64_t ntpSeconds = static_cast<uint64_t>(unixTime) + NTP_EPOCH_OFFSET;

    NTPPacket response{};
    response.recvTime = ntpSeconds;
    response.xmitTime = ntpSeconds;

    SOCKET sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sock == INVALID_SOCKET) 
    {
        WSACleanup();
        return 1;
    }

    sockaddr_in addr{};
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = INADDR_ANY;
    addr.sin_port = htons(123);

    if (bind(sock, reinterpret_cast<SOCKADDR*>(&addr), sizeof(addr)) == SOCKET_ERROR) 
    {
        closesocket(sock);
        WSACleanup();
        return 1;
    }

    cout << "NTP server running on port 123..." << endl;

    sockaddr_in clientAddr{};
    int clientAddrLen = sizeof(clientAddr);
    char buf[NTP_PACKET_SIZE];

    for (;;) {
        int recvLen = recvfrom(sock, buf, NTP_PACKET_SIZE, 0,
            reinterpret_cast<SOCKADDR*>(&clientAddr), &clientAddrLen);
        if (recvLen == SOCKET_ERROR) 
        {
            if (WSAGetLastError() != WSAECONNRESET) cerr << "recvfrom failed" << endl;
            continue;
        }

        NTPPacket* req = reinterpret_cast<NTPPacket*>(buf);
        if (((req->liVnMode >> 5) & 0x7) != NTP_HEADER_MODE_CLIENT ||
            ((req->liVnMode >> 3) & 0x7) != NTP_HEADER_VERSION) continue;

        response.origTime = *reinterpret_cast<uint64_t*>(req->origTime);
        if (sendto(sock, reinterpret_cast<char*>(&response), NTP_PACKET_SIZE, 0,
            reinterpret_cast<SOCKADDR*>(&clientAddr), clientAddrLen) == SOCKET_ERROR) 
        {
            cerr << "sendto failed" << endl;
        }
    }

    closesocket(sock);
    WSACleanup();
    return 0;
}